import numpy as np
import random
import torch
from torch.nn import Module, Parameter
from torch.nn import functional as F
import network
from performance_evaluation import group_reward
import network_topology as nt
import task_generator
import networkx as nx
from const import TASK_NUM, GROUP_SIZE, EPISODE_SIZE
import os


def reward_evaluator(G):
    edge_num = G.edges.__len__()
    node_num = G.nodes.__len__()
    # GAT = network.GAT_with_fc(TASK_NUM, edge_num, 1, node_num, num_h=4, dropout=True, alpha=0.1)
    GAT = network.GAT_with_fc(TASK_NUM, edge_num, 1, node_num, 1, dropout=True, alpha=0.1, num_h=2)
    return GAT


def weight_update_sim(G):
    new_weight = np.empty(26)
    for i in range(0, 17):
        new_weight[i] = random.randint(5, 20)
    G = nt.weight_update(G, new_weight)

    adj = np.zeros((17, 17))

    edge_labels = nx.get_edge_attributes(G, 'weight')
    for key, val in edge_labels.items():
        a = key[0]
        b = key[1]
        adj[a][b] = val
        adj[b][a] = val

    # adj = nx.adjacency_matrix(G).todense()

    return G, adj


'''
训练过程
'''

# 网络基础性质生成
G, _ = nt.GBN()
capacity_matrix = nt.generate_capacity_matrix(G, 100, 150)
v_matrix = nt.generate_v_matrix(G, 10, 30)

# 神经网络，优化器，任务生成器初始化，并表明训练区域
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
GAT = reward_evaluator(G).to(device)
optimizer = torch.optim.Adam(GAT.parameters(), lr=0.01)
generator = task_generator.taskGenerator(G)
loss_avg = []
print(GAT)
GAT.train()

for epidode in range(EPISODE_SIZE):

    task_group = generator.generate_group(GROUP_SIZE, 50, 100)
    G, weighted_adj = weight_update_sim(G)
    reward_list = group_reward(G, task_group, capacity_matrix, v_matrix)[1:-1]
    reward_list = torch.tensor(reward_list, dtype=torch.float32).to(device)
    loss_list = torch.empty_like(reward_list)

    for i in range(0, GROUP_SIZE - 1):
        optimizer.zero_grad()
        task_matrix = torch.tensor(task_group[i]['task_lists']['task_matrix'], dtype=torch.float32).to(device)
        weighted_adj = torch.tensor(weighted_adj, dtype=torch.float32).to(device)
        output = GAT.forward(task_matrix.T, weighted_adj)  # 这里要的是node_num * input_num形状的特征矩阵
        loss = F.mse_loss(output, reward_list[i - 1])
        loss_list[i - 1] = output
        loss.backward()
        optimizer.step()

    if epidode % 200 == 0:
        print(loss_list - reward_list)
        print('episode ' + str(epidode))
        avg_loss = F.mse_loss(loss_list, reward_list)
        print(avg_loss)
        loss_avg.append(avg_loss)

    if epidode % 1000 == 0:
        print(loss_avg)

torch.save(GAT, 'model.pth')
